package org.xyl.service;

import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.xyl.engine.ChromaInitializer;
import org.xyl.service.VectorLibraryServiceInterface;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class VectorLibraryService implements VectorLibraryServiceInterface {

    @Autowired
    private EmbeddingStore<TextSegment> embeddingStore;

    @Autowired
    private EmbeddingModel embeddingModel;

    // 向量库元数据管理
    private final Map<String, VectorLibraryMetadata> vectorLibraries = new ConcurrentHashMap<>();

    /**
     * 创建新的向量库
     * @param libraryName 向量库名称
     * @return 是否创建成功
     */
    @Override
    public boolean createVectorLibrary(String libraryName) {
        if (vectorLibraries.containsKey(libraryName)) {
            return false; // 库已存在
        }
        vectorLibraries.put(libraryName, new VectorLibraryMetadata(libraryName));
        System.out.println("创建向量库: " + libraryName);
        return true;
    }

    /**
     * 删除向量库
     * @param libraryName 向量库名称
     * @return 是否删除成功
     */
    @Override
    public boolean deleteVectorLibrary(String libraryName) {
        if (!vectorLibraries.containsKey(libraryName)) {
            return false; // 库不存在
        }
        vectorLibraries.remove(libraryName);
        // 注意：这里只是删除元数据，实际的向量数据需要根据存储类型进行清理
        System.out.println("删除向量库: " + libraryName);
        return true;
    }

    /**
     * 获取所有向量库列表
     * @return 向量库名称列表
     */
    @Override
    public List<String> listVectorLibraries() {
        return List.copyOf(vectorLibraries.keySet());
    }

    /**
     * 内部使用的添加文档方法（String元数据）
     * @param libraryName 向量库名称
     * @param documentText 文档文本
     * @param documentMetadata 文档元数据
     * @return 是否添加成功
     */
    public boolean addDocumentWithStringMetadata(String libraryName, String documentText, Map<String, String> documentMetadata) {
        if (!vectorLibraries.containsKey(libraryName)) {
            return false; // 库不存在
        }

        try {
            // 为文档文本创建嵌入
            Embedding embedding = embeddingModel.embed(documentText).content();
            
            // 创建文本段
            TextSegment segment = TextSegment.from(documentText);
            
            // 添加到向量存储
            embeddingStore.add(embedding, segment);
            
            // 更新库元数据
            VectorLibraryMetadata libraryMetadata = vectorLibraries.get(libraryName);
            libraryMetadata.incrementDocumentCount();
            
            return true;
        } catch (Exception e) {
            System.err.println("添加文档到向量库失败: " + e.getMessage());
            return false;
        }
    }

    /**
     * 从指定向量库中检索相关文档
     * @param libraryName 向量库名称
     * @param query 查询文本
     * @param maxResults 返回结果数量
     * @return 相关文档匹配列表
     */
    @Override
    public List<EmbeddingMatch<TextSegment>> searchInLibrary(String libraryName, String query, int maxResults) {
        if (!vectorLibraries.containsKey(libraryName)) {
            return List.of(); // 库不存在，返回空列表
        }

        try {
            // 为查询创建嵌入
            Embedding queryEmbedding = embeddingModel.embed(query).content();
            
            // 从向量存储中查找相关内容
            return embeddingStore.findRelevant(queryEmbedding, maxResults);
        } catch (Exception e) {
            System.err.println("在向量库中搜索失败: " + e.getMessage());
            return List.of();
        }
    }

    /**
     * 获取向量库统计信息
     * @param libraryName 向量库名称
     * @return 库元数据
     */
    public VectorLibraryMetadata getLibraryMetadata(String libraryName) {
        return vectorLibraries.get(libraryName);
    }
    
    @Override
    public List<EmbeddingMatch<TextSegment>> searchInVectorLibrary(String libraryName, Embedding queryEmbedding, int topK) {
        if (!vectorLibraries.containsKey(libraryName)) {
            return List.of(); // 库不存在，返回空列表
        }

        try {
            // 从向量存储中查找相关内容
            return embeddingStore.findRelevant(queryEmbedding, topK);
        } catch (Exception e) {
            System.err.println("在向量库中搜索失败: " + e.getMessage());
            return List.of();
        }
    }
    
    @Override
    public void addDocumentToLibrary(String libraryName, String document, Map<String, Object> metadata) {
        try {
            // 创建文本片段（直接使用文本内容，不使用metadata参数）
            TextSegment textSegment = TextSegment.from(document);
            
            // 生成嵌入向量
            Embedding embedding = embeddingModel.embed(textSegment).content();
            
            // 存储到嵌入存储
            embeddingStore.add(embedding, textSegment);
            
            // 更新库元数据
            if (vectorLibraries.containsKey(libraryName)) {
                VectorLibraryMetadata libraryMetadata = vectorLibraries.get(libraryName);
                libraryMetadata.incrementDocumentCount();
            }
        } catch (Exception e) {
            System.err.println("添加文档到向量库失败: " + e.getMessage());
        }
    }
    
    @Override
    public Map<String, Object> getVectorLibraryStats(String libraryName) {
        VectorLibraryMetadata metadata = getLibraryMetadata(libraryName);
        if (metadata == null) {
            return Map.of();
        }
        return Map.of(
            "name", metadata.getName(),
            "documentCount", metadata.getDocumentCount(),
            "createdAt", metadata.getCreatedAt(),
            "updatedAt", metadata.getUpdatedAt()
        );
    }

    // 向量库元数据类
    public static class VectorLibraryMetadata {
        private String name;
        private int documentCount = 0;
        private long createdAt = System.currentTimeMillis();
        private long updatedAt = System.currentTimeMillis();

        public VectorLibraryMetadata(String name) {
            this.name = name;
        }

        public String getName() {
            return name;
        }

        public int getDocumentCount() {
            return documentCount;
        }

        public long getCreatedAt() {
            return createdAt;
        }

        public long getUpdatedAt() {
            return updatedAt;
        }

        public void incrementDocumentCount() {
            this.documentCount++;
            this.updatedAt = System.currentTimeMillis();
        }
    }
}